#encoding: utf-8
import xml.sax.handler

class KnowledgeBaseHandler(xml.sax.handler.ContentHandler):
    '''This handler is used to parse the query XML and then extract
       all doc id of Chinese query which is stored in setting.
    '''
    def __init__(self):
        #
        self.link_id = ""
        self.buffer = ""
        self.in_record = False
        self.entity_dict = {}
        self.entity_title_dict = {}

    def startElement(self, name, attributes):
        if name == "entity":
            entity_id = attributes['id']
            entity_name = attributes['name']
            entity_title = attributes['wiki_title']
            if entity_id not in self.entity_dict:
                self.entity_dict[entity_id] = [entity_name]
            else:
                if entity_name not in self.entity_dict[entity_id]:
                    self.entity_dict[entity_id].append(entity_name)

            if entity_id not in self.entity_title_dict:
                self.entity_title_dict[entity_id] = entity_title
            else:
                if entity_title not in self.entity_title_dict[entity_id]:
                    self.entity_title_dict[entity_id].append(entity_title)
                print('The same entity id appears more than one time!')
                print('Detail:\n\tentity_id:{0}\n\tentity_title:{1}'.\
                        format(self.entity_title_dict[entity_id]))
        if name == "link" and 'entity_id' in attributes:
            entity_id = attributes['entity_id']
            self.link_id = entity_id
            self.buffer = ""
            self.in_record = True

    def characters(self, data):
        if self.in_record:
            self.buffer += data

    def endElement(self, name):
        if name == "link" and self.in_record:
            self.buffer = self.buffer.strip()
            if self.link_id not in self.entity_dict:
                self.entity_dict[self.link_id] = [self.buffer]
            else:
                if self.buffer not in self.entity_dict[self.link_id]:
                    self.entity_dict[self.link_id].append(self.buffer)
            self.link_id = ""
            self.buffer = ""
            self.in_record = False
